package com.test.ai.dao;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JavaType;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.time.Instant;
import java.util.Collections;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;

/**
 * Created with IntelliJ IDEA.
 *
 * @description:
 * @author: liuziyang
 * @since: 8/21/25 5:01 PM
 * @modifiedBy:
 * @version: 1.0
 */
public class RedisChatMemory implements ChatMemory {
  private static final String JSON_ARRAY_STARTER = "[";
  private static final String CHAT_MEMORY_KEY_PREFIX = "chat:memory:";
  private final RedisTemplate<String, String> redisTemplate;
  private final ObjectMapper objectMapper;
  private final long messageTtl;

  public RedisChatMemory(RedisTemplate<String, String> redisTemplate, ObjectMapper objectMapper) {
    // 默认30天
    this(redisTemplate, objectMapper, 30 * 24 * 60 * 60);
  }

  public RedisChatMemory(
      RedisTemplate<String, String> redisTemplate, ObjectMapper objectMapper, long messageTtl) {
    this.redisTemplate = redisTemplate;
    this.objectMapper = objectMapper;
    this.messageTtl = messageTtl;
  }

  @Override
  public void add(String conversationId, List<Message> messages) {
    if (messages == null || messages.isEmpty()) {
      return;
    }

    String key = CHAT_MEMORY_KEY_PREFIX + conversationId;
    ZSetOperations<String, String> zSetOps = redisTemplate.opsForZSet();
    // 用时间戳作为score
    long baseTimestamp = Instant.now().toEpochMilli();

    // 为每条消息生成递增的时间戳，确保消息顺序
    for (int i = 0; i < messages.size(); i++) {
      try {
        String messageJson = objectMapper.writeValueAsString(messages.get(i));
        // 使用baseTimestamp + i确保同批次消息的顺序
        zSetOps.add(key, messageJson, baseTimestamp + i);
      } catch (JsonProcessingException e) {
        throw new RuntimeException("消息序列化失败", e);
      }
    }

    redisTemplate.expire(key, messageTtl, java.util.concurrent.TimeUnit.SECONDS);
  }

  @Override
  public List<Message> get(String conversationId, int lastN) {
    if (lastN <= 0) {
      return Collections.emptyList();
    }

    String key = CHAT_MEMORY_KEY_PREFIX + conversationId;
    ZSetOperations<String, String> zSetOps = redisTemplate.opsForZSet();
    Set<String> messageJsons = zSetOps.range(key, -lastN, -1);

    if (messageJsons == null || messageJsons.isEmpty()) {
      return Collections.emptyList();
    }

    return messageJsons.stream().map(this::deserializeMessage).collect(Collectors.toList());
  }

  @Override
  public void clear(String conversationId) {
    redisTemplate.delete(CHAT_MEMORY_KEY_PREFIX + conversationId);
  }

  private Message deserializeMessage(String json) {
    try {
      if (json.trim().startsWith(JSON_ARRAY_STARTER)) {
        JavaType type =
            objectMapper.getTypeFactory().constructParametricType(List.class, Message.class);
        List<Message> messages = objectMapper.readValue(json, type);
        return messages.isEmpty() ? null : messages.getFirst();
      }
      // 使用Message混合类型反序列化
      return objectMapper.readValue(json, Message.class);
    } catch (JsonProcessingException e) {
      throw new RuntimeException("消息反序列化失败: " + json, e);
    }
  }
}
